//
// Created by ZhaoXiaoFei on 2022/8/5.
//

#include "pose_graph.hpp"

PoseGraph::PoseGraph() {
    posegraph_visualization = new CameraPoseVisualization(1.0, 0.0, 1.0, 1.0);
    posegraph_visualization->setScale(0.1);
    posegraph_visualization->setLineWidth(0.01);

    r_drift = Eigen::Matrix3d::Identity();
    yaw_drift = 0;
    t_drift = Eigen::Vector3d::Zero();

    t_optimization = std::thread(&PoseGraph::optimize4DOF, this);

    w_r_vio = Eigen::Matrix3d::Identity();
    w_t_vio = Eigen::Vector3d::Zero();
    sequence_cnt = 0;
    sequence_loop.push_back(0);
    global_index = 0;
    earliest_loop_index = -1;
}
int PoseGraph::detectLoop(KeyFrame* cur_kf, int cur_kf_index){
//    LOG(INFO) << "11111111111111" <<  std::endl;
    cv::Mat compressed_image;
    if(DEBUG_IMAGE){
        int feature_num = cur_kf->keyPoints.size();
        cv::resize(cur_kf->image, compressed_image, cv::Size(376, 240));
        cv::putText(compressed_image, "feature_num:" + to_string(feature_num), cv::Point2f(10,10),cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(255));
        image_pool[cur_kf_index] = compressed_image;
    }
    DBoW2::QueryResults ret;
//    LOG(INFO) << "11111111111111" <<  std::endl;
    bd.query(cur_kf->briefDescriptors, ret, 4, cur_kf_index - 50);
    LOG(INFO) << "Searching for Image " << cur_kf_index << "  " << ret << endl;
    bd.add(cur_kf->briefDescriptors);
    bool find_loop = false;
//    LOG(INFO) << "11111111111111" <<  std::endl;
    cv::Mat loop_result;
    if (DEBUG_IMAGE){
        loop_result = compressed_image.clone();
        if (ret.size() > 0){
            putText(loop_result, "neighbour score:" + to_string(ret[0].Score), cv::Point2f(10, 50), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255));
        }
    }
    if (DEBUG_IMAGE){
        for (int i = 0; i < ret.size(); i++){
            int tmp_index = ret[i].Id;
            auto it = image_pool.find(tmp_index);
            cv::Mat tmp_image = (it->second).clone();
            putText(tmp_image, "index:  " + to_string(tmp_index) + "loop score:" + to_string(ret[i].Score), cv::Point2f(10, 50), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255));
            cv::hconcat(loop_result, tmp_image, loop_result);
        }
    }
    if (DEBUG_IMAGE)
    {
//        cv::imshow("loop_result", loop_result);
//        cv::waitKey(20);
    }

    if(ret.size() >= 1 && ret[0].Score > 0.05){
        for(int i = 1; i < ret.size(); i++){
            if(ret[i].Score  > 0.015){
                find_loop = true;
            }
        }
    }
    if(find_loop && cur_kf_index > 50){
        int min_index = -1;
        for(int i = 0; i < ret.size(); i++){
            if(min_index == -1 || (ret[i].Id < min_index && ret[i].Score > 0.015)){
                min_index = ret[i].Id;
            }
        }
        return min_index;
    }
    else{
        return -1;
    }
}

void PoseGraph::registerPub(ros::NodeHandle &n) {
    pub_pg_path = n.advertise<nav_msgs::Path>("pose_graph_path", 1000);
    pub_base_path = n.advertise<nav_msgs::Path>("base_path", 1000);
    pub_pose_graph = n.advertise<visualization_msgs::MarkerArray>("pose_graph", 1000);
    for (int i = 1; i < 10; i++)
        pub_path[i] = n.advertise<nav_msgs::Path>("path_" + to_string(i), 1000);
}

void PoseGraph::addKeyFrameIntoVoc(KeyFrame *cur_kf) {
    cv::Mat compressed_image;
    if (DEBUG_IMAGE){
        int feature_num = cur_kf->keyPoints.size();
        cv::resize(cur_kf->image, compressed_image, cv::Size(376, 240));
        putText(compressed_image, "feature_num:" + to_string(feature_num), cv::Point2f(10, 10), cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(255));
        image_pool[cur_kf->index] = compressed_image;
    }
    bd.add(cur_kf->briefDescriptors);
}

void PoseGraph::publish() {
    for(int i = 1; i <= sequence_cnt; i++){
        pub_pg_path.publish(path[i]);
        pub_path[i].publish(path[i]);
        posegraph_visualization->publish_by(pub_pose_graph, path[sequence_cnt].header);
    }
    base_path.header.frame_id = "world";
    pub_base_path.publish(base_path);
}

void PoseGraph::loadVocabulary(std::string voc_path) {
    voc = new BriefVocabulary(voc_path);
    bd.setVocabulary(*voc, false, 0);
}

void PoseGraph::addKeyFrame(KeyFrame *cur_kf, bool flag_detect_loop) {
//    LOG(INFO) << "11111111111111" <<  std::endl;
    if(sequence_cnt != cur_kf->sequence){
        sequence_cnt++;
        sequence_loop.push_back(0);
        w_r_vio = Eigen::Matrix3d::Identity();
        w_t_vio = Eigen::Vector3d::Zero();
        m_drift.lock();
        r_drift = Eigen::Matrix3d::Identity();
        t_drift = Eigen::Vector3d::Zero();
        m_drift.unlock();
    }
    Eigen::Vector3d vio_t_cur;
    Eigen::Matrix3d vio_R_cur;
    cur_kf->getVioPose(vio_R_cur, vio_t_cur);
    vio_R_cur = w_r_vio * vio_R_cur;
    vio_t_cur = w_r_vio * vio_t_cur + w_t_vio;
    cur_kf->updateVioPose(vio_R_cur, vio_t_cur);
    cur_kf->index = global_index++;
    int loop_index = -1;
//    LOG(INFO) << "11111111111111" <<  std::endl;
    if(flag_detect_loop){
        loop_index = detectLoop(cur_kf, cur_kf->index);
    }
    else{
        addKeyFrameIntoVoc(cur_kf);
    }
    if(loop_index != -1){
        KeyFrame* old_kf = getKeyFrame(loop_index);

        if(cur_kf->findConnection(old_kf)){
            if(earliest_loop_index == -1 || earliest_loop_index > loop_index){
                earliest_loop_index = loop_index;
            }
            //            R_w1_bi, R_w1_bi_bj, R_w1^_bj; ^代表squence引起的问题;
            Eigen::Matrix3d w_R_old, w_R_cur, vio_R_cur;
            Eigen::Vector3d w_T_old, w_T_cur, vio_T_cur;
            //不应该包括drift;
            old_kf->getVioPose(w_R_old, w_T_old);
            cur_kf->getVioPose(vio_R_cur, vio_T_cur);

            Eigen::Matrix3d relativeR; Eigen::Vector3d relativeT;
            relativeR = cur_kf->getLoopRelativeR();
            relativeT = cur_kf->getLoopRelativeT();
            //R_w1_bi -> R_w1_bj;
            w_R_cur = w_R_old * relativeR;
            w_T_cur = w_R_old * relativeT + w_T_old;
            double shift_yaw;
            Eigen::Matrix3d shift_r;
            Eigen::Vector3d shift_t;
//            LOG(INFO) << "111111111111111 " << std::endl;
            shift_yaw = Utility::R2ypr(w_R_cur).x() - Utility::R2ypr(vio_R_cur).x();
            shift_r = Utility::ypr2R(Eigen::Vector3d(shift_yaw, 0, 0));
            shift_t = w_T_cur - w_R_cur * vio_R_cur.transpose() * vio_T_cur;
            if(old_kf->sequence != cur_kf->sequence && sequence_loop[cur_kf->sequence] == 0){
                w_r_vio = shift_r;
                w_t_vio = shift_t;
                vio_R_cur = w_r_vio * vio_R_cur;
                vio_T_cur = w_r_vio * vio_t_cur + w_t_vio;
                cur_kf->updateVioPose(vio_R_cur, vio_T_cur);
                std::list<KeyFrame*>::iterator it = keyFrameList.begin();
                for(; it != keyFrameList.end(); it++){
                    if((*it)->sequence == cur_kf->sequence){
                        Eigen::Matrix3d R_vio;
                        Eigen::Vector3d T_vio;
                        (*it)->getVioPose(R_vio, T_vio);
                        R_vio = w_r_vio * R_vio;
                        T_vio = w_r_vio * T_vio + w_t_vio;
                        (*it)->updateVioPose(R_vio, T_vio);
                    }
                }
                sequence_loop[cur_kf->sequence] = 1;
            }
            m_optimize_buf.lock();
            optimize_buf.push(cur_kf->index);
            m_optimize_buf.unlock();
        }
    }
    m_keyframelist.lock();
    Eigen::Matrix3d R;
    Eigen::Vector3d T;
    cur_kf->getVioPose(R, T);
    R = r_drift * R;
    T = r_drift * T + t_drift;
    cur_kf->updatePose(R, T);
    Eigen::Quaterniond tempQ(R);
    geometry_msgs::PoseStamped poseStamped;
    poseStamped.header.stamp = ros::Time(cur_kf->time_stamp);
    poseStamped.header.frame_id = "world";
    poseStamped.pose.position.x = T.x() + VISUALIZATION_SHIFT_X;
    poseStamped.pose.position.y = T.y() + VISUALIZATION_SHIFT_Y;
    poseStamped.pose.position.z = T.z();
    poseStamped.pose.orientation.w = tempQ.w();
    poseStamped.pose.orientation.x = tempQ.x();
    poseStamped.pose.orientation.y = tempQ.y();
    poseStamped.pose.orientation.z = tempQ.z();

    path[sequence_cnt].header = poseStamped.header;
    path[sequence_cnt].poses.push_back(poseStamped);

    keyFrameList.push_back(cur_kf);
    publish();
    m_keyframelist.unlock();
    LOG(INFO) << "addKeyFrame end " << std::endl;
}

KeyFrame *PoseGraph::getKeyFrame(int loopIndex) {

    std::list<KeyFrame*>::iterator it;
    for(it = keyFrameList.begin(); it != keyFrameList.end(); it++){
        if((*it)->index == loopIndex){
            break;
        }
    }
    return it != keyFrameList.end() ? *it : nullptr;
}

void PoseGraph::optimize4DOF() {
    while(true){

        int cur_index = -1;
        int first_loop_index = -1;
        m_optimize_buf.lock();
        while(!optimize_buf.empty()){
            cur_index = optimize_buf.front();
            first_loop_index = earliest_loop_index;
            optimize_buf.pop();
        }
        m_optimize_buf.unlock();
        if(cur_index != -1){
            m_keyframelist.lock();
            LOG(INFO) << "optimize4DOF" << std::endl;
            KeyFrame* cur_keyFrame = getKeyFrame(cur_index);

            int max_length = cur_index + 1;
            double para_T[max_length][3];
            Eigen::Quaterniond Q[max_length];
            double para_euler[max_length][3];
            double sequence_array[max_length];
//            LOG(INFO) << "optimize4DOF" << std::endl;
            ceres::Problem problem;
            ceres::Solver::Options options;
            options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;
            options.max_num_iterations = 5;
            ceres::Solver::Summary summary;
            ceres::LossFunction *lossFunction = new ceres::HuberLoss(0.1);
            ceres::LocalParameterization* angleLocalParameterization = AngleLocalParameterization::Create();
//            LOG(INFO) << "optimize4DOF" << std::endl;
            int i = 0;
            std::list<KeyFrame*>::iterator iter;
            for(iter = keyFrameList.begin() ; iter != keyFrameList.end(); iter++){
//                LOG(INFO) << "optimize4DOF" << std::endl;
                if((*iter)->index < first_loop_index){
                    continue;
                }
                (*iter)->local_index = i;
                Eigen::Matrix3d temp_r;
                Eigen::Vector3d temp_t;
                (*iter)->getVioPose(temp_r, temp_t);
//                LOG(INFO) << "optimize4DOF" << std::endl;
                para_T[i][0] = temp_t.x();
                para_T[i][1] = temp_t.y();
                para_T[i][2] = temp_t.z();
//                LOG(INFO) << "optimize4DOF" << std::endl;
                Q[i] = Eigen::Quaterniond(temp_r);
                Eigen::Vector3d temp_euler = Utility::R2ypr(temp_r);
//                LOG(INFO) << "optimize4DOF" << std::endl;
                para_euler[i][0] = temp_euler.x();
                para_euler[i][1] = temp_euler.y();
                para_euler[i][2] = temp_euler.z();
//                LOG(INFO) << "optimize4DOF" << std::endl;
                sequence_array[i] = (*iter)->sequence;

                problem.AddParameterBlock(para_T[i], 3);
                problem.AddParameterBlock(para_euler[i], 1, angleLocalParameterization);
//                LOG(INFO) << "optimize4DOF" << std::endl;
                if((*iter)->index == first_loop_index || (*iter)->sequence == 0){
                    problem.SetParameterBlockConstant(para_euler[i]);
                    problem.SetParameterBlockConstant(para_T[i]);
                }
//                LOG(INFO) << "optimize4DOF" << std::endl;
                // 序列边
                for(int j = 1; j < 5; j++){
                    if(i - j >= 0 && sequence_array[i] == sequence_array[i - j]){
                        Eigen::Vector3d euler_i = Utility::R2ypr(Q[i - j].toRotationMatrix());
                        //此处的相对值 属于预测值
                        //优化过程，残差 = 观测 - 状态  而状态这里取得是预测值；
                        Eigen::Vector3d reltiveT = Q[i - j] .inverse() * Eigen::Vector3d(para_T[i][0] - para_T[i - j][0], para_T[i][1] - para_T[i - j][1], para_T[i][2] - para_T[i - j][2]);
                        double relativeYaw = para_euler[i][0] - para_euler[i - j][0];
                        ceres::CostFunction* costFunction = FourDofError::Create(relativeYaw, euler_i.y(), euler_i.z(), reltiveT.x(), reltiveT.y(), reltiveT.z());
                        problem.AddResidualBlock(costFunction, nullptr, para_euler[i-j], para_T[i-j], para_euler[i], para_T[i]);
                    }
                }
//                LOG(INFO) << "optimize4DOF" << std::endl;
                //回环边
                if((*iter)->has_loop){
                    assert((*iter)->loop_index >= first_loop_index);
                    int connect_i = getKeyFrame((*iter)->loop_index)->local_index;
                    Eigen::Vector3d euler_i = Utility::R2ypr(Q[connect_i].toRotationMatrix());
                    Eigen::Vector3d reltiveT = (*iter)->getLoopRelativeT();
                    double relativeYaw = (*iter)->getLoopRelativeYaw();
                    ceres::CostFunction* costFunction = FourDofError::Create(relativeYaw, euler_i.y(), euler_i.z(), reltiveT.x(), reltiveT.y(), reltiveT.z());
                    problem.AddResidualBlock(costFunction, nullptr, para_euler[connect_i], para_T[connect_i], para_euler[i], para_T[i]);
                }
//                LOG(INFO) << "optimize4DOF" << std::endl;
                if((*iter)->index == cur_index){
                    break;
                }
                i++;
            }
            m_keyframelist.unlock();

            ceres::Solve(options, &problem, &summary);
            LOG(INFO) << summary.BriefReport() << "\n";

            for (int j = 0 ; j < i; j++)
            {
                printf("optimize i: %d p: %f, %f, %f\n", j, para_T[j][0], para_T[j][1], para_T[j][2]);
            }

            m_keyframelist.lock();
            i = 0;
            for(iter = keyFrameList.begin(); iter != keyFrameList.end(); iter++){
                if((*iter)->index < first_loop_index){
                    continue;
                }

                Eigen::Matrix3d temp_r = Utility::ypr2R(Eigen::Vector3d(para_euler[i][0], para_euler[i][1],para_euler[i][2]));
                Eigen::Vector3d temp_t = Eigen::Vector3d(para_T[i][0], para_T[i][1],para_T[i][2]);

                (*iter)->updatePose(temp_r, temp_t);

                if((*iter)->index == cur_index){
                    break;
                }
                i++;
            }
            Eigen::Vector3d w_t, vio_t;
            Eigen::Matrix3d w_r, vio_r;
            cur_keyFrame->getVioPose(vio_r, vio_t);
            cur_keyFrame->getPose(w_r, w_t);
            m_drift.lock();
            yaw_drift = Utility::R2ypr(w_r).x() - Utility::R2ypr(vio_r).x();
            r_drift = Utility::ypr2R(Eigen::Vector3d(yaw_drift, 0, 0));
            t_drift = w_t - r_drift * vio_t;
            m_drift.unlock();
            iter++;
            for( ; iter != keyFrameList.end(); iter++){
                Eigen::Vector3d t;
                Eigen::Matrix3d r;
                (*iter)->getVioPose(r, t);
                r = r_drift * r;
                t = t_drift + r_drift * t;
                (*iter)->updatePose(r, t);
            }
            m_keyframelist.unlock();
            updatePath();
        }
        std::chrono::milliseconds duration(2000);
        std::this_thread::sleep_for(duration);
    }
}

void PoseGraph::updatePath()
{
    m_keyframelist.lock();
    list<KeyFrame*>::iterator it;
    for (int i = 1; i <= sequence_cnt; i++)
    {
        path[i].poses.clear();
    }
    base_path.poses.clear();
    posegraph_visualization->reset();

    for (it = keyFrameList.begin(); it != keyFrameList.end(); it++)
    {
        Eigen::Vector3d P;
        Eigen::Matrix3d R;
        (*it)->getPose(R, P);
        Eigen::Quaterniond Q;
        Q = R;
//        printf("path p: %f, %f, %f\n",  P.x(),  P.z(),  P.y() );

        geometry_msgs::PoseStamped pose_stamped;
        pose_stamped.header.stamp = ros::Time((*it)->time_stamp);
        pose_stamped.header.frame_id = "world";
        pose_stamped.pose.position.x = P.x() + VISUALIZATION_SHIFT_X;
        pose_stamped.pose.position.y = P.y() + VISUALIZATION_SHIFT_Y;
        pose_stamped.pose.position.z = P.z();
        pose_stamped.pose.orientation.x = Q.x();
        pose_stamped.pose.orientation.y = Q.y();
        pose_stamped.pose.orientation.z = Q.z();
        pose_stamped.pose.orientation.w = Q.w();
        if((*it)->sequence == 0)
        {
            base_path.poses.push_back(pose_stamped);
            base_path.header = pose_stamped.header;
        }
        else
        {
            path[(*it)->sequence].poses.push_back(pose_stamped);
            path[(*it)->sequence].header = pose_stamped.header;
        }


        //draw local connection
        if (0)
        {
            list<KeyFrame*>::reverse_iterator rit = keyFrameList.rbegin();
            list<KeyFrame*>::reverse_iterator lrit;
            for (; rit != keyFrameList.rend(); rit++)
            {
                if ((*rit)->index == (*it)->index)
                {
                    lrit = rit;
                    lrit++;
                    for (int i = 0; i < 4; i++)
                    {
                        if (lrit == keyFrameList.rend())
                            break;
                        if((*lrit)->sequence == (*it)->sequence)
                        {
                            Eigen::Vector3d conncected_P;
                            Eigen::Matrix3d connected_R;
                            (*lrit)->getPose(connected_R, conncected_P);
                            posegraph_visualization->add_edge(P, conncected_P);
                        }
                        lrit++;
                    }
                    break;
                }
            }
        }
        if (1)
        {
            if ((*it)->has_loop && (*it)->sequence == sequence_cnt)
            {

                KeyFrame* connected_KF = getKeyFrame((*it)->loop_index);
                Eigen::Vector3d connected_P;
                Eigen::Matrix3d connected_R;
                connected_KF->getPose(connected_R, connected_P);
                //(*it)->getVioPose(P, R);
                (*it)->getPose(R, P);
                if((*it)->sequence > 0)
                {
                    posegraph_visualization->add_loopedge(P, connected_P + Eigen::Vector3d(VISUALIZATION_SHIFT_X, VISUALIZATION_SHIFT_Y, 0));
                }
            }
        }
    }
    publish();
    m_keyframelist.unlock();
}

void PoseGraph::updateKeyFrameLoop(int index, Eigen::Matrix<double, 8, 1> &_loop_info) {
    KeyFrame* keyFrame = getKeyFrame(index);
    keyFrame->updateLoop(_loop_info);
    if (abs(_loop_info(7)) < 30.0 && Eigen::Vector3d(_loop_info(0), _loop_info(1), _loop_info(2)).norm() < 20.0){
        if(FAST_RELOCALIZATION){
            KeyFrame* old_kf = getKeyFrame(keyFrame->loop_index);
            Eigen::Vector3d w_P_old, w_P_cur, vio_P_cur;
            Eigen::Matrix3d w_R_old, w_R_cur, vio_R_cur;
            old_kf->getPose(w_R_old, w_P_old);
            keyFrame->getVioPose(vio_R_cur, vio_P_cur);

            Eigen::Vector3d relative_t;
            Eigen::Quaterniond relative_q;
            relative_t = keyFrame->getLoopRelativeT();
            relative_q = (keyFrame->getLoopRelativeR());
            w_P_cur = w_R_old * relative_t + w_P_old;
            w_R_cur = w_R_old * relative_q;
            double shift_yaw;
            Eigen::Matrix3d shift_r;
            Eigen::Vector3d shift_t;
            shift_yaw = Utility::R2ypr(w_R_cur).x() - Utility::R2ypr(vio_R_cur).x();
            shift_r = Utility::ypr2R(Eigen::Vector3d(shift_yaw, 0, 0));
            shift_t = w_P_cur - w_R_cur * vio_R_cur.transpose() * vio_P_cur;

            m_drift.lock();
            yaw_drift = shift_yaw;
            r_drift = shift_r;
            t_drift = shift_t;
            m_drift.unlock();
        }
    }

}